# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Handles directives.

This converter removes the directive functions from the code and moves the
information they specify into AST annotations. It is a specialized form of
static analysis, one that is specific to AutoGraph.

Note that this requires that the actual directive functions are static - that
is, they do not change at runtime. So if you do something like this:

  tf.autograph.set_loop_options = <new function>

Then the directive will may no longer be recognized. Furthermore, if the
converted function is cached, such an action may be irreversible.
"""

import inspect

import gast

from nvidia.dali._autograph.core import converter
from nvidia.dali._autograph.lang import directives
from nvidia.dali._autograph.pyct import anno


STATIC_VALUE = "static_value"
"""Used for AST annotations, see visit_Name."""


class _LoopScope(object):
    def __init__(self):
        self.ast_node = None
        self.statements_visited = 0


def _map_args(call_node, function):
    """Maps AST call nodes to the actual function's arguments.

    Args:
      call_node: ast.Call
      function: Callable[..., Any], the actual function matching call_node
    Returns:
      Dict[Text, ast.AST], mapping each of the function's argument names to
      the respective AST node.
    Raises:
        ValueError: if the default arguments are not correctly set
    """
    args = call_node.args
    kwds = {kwd.arg: kwd.value for kwd in call_node.keywords}
    call_args = inspect.getcallargs(function, *args, **kwds)

    # Keyword arguments not specified in kwds will be mapped to their defaults,
    # which are Python values. Since we don't currently have a way to transform
    # those into AST references, we simply remove them. By convention, directives
    # use UNSPECIFIED as default value for optional arguments. No other
    # defaults should be present.
    unexpected_defaults = []
    for k in call_args:
        if (
            k not in kwds
            and call_args[k] not in args
            and call_args[k] is not directives.UNSPECIFIED
        ):
            unexpected_defaults.append(k)
    if unexpected_defaults:
        raise ValueError(
            "Unexpected keyword argument values, %s, for function %s"
            % (zip(unexpected_defaults, [call_args[k] for k in unexpected_defaults]), function)
        )
    return {k: v for k, v in call_args.items() if v is not directives.UNSPECIFIED}


class DirectivesTransformer(converter.Base):
    """Parses compiler directives and converts them into AST annotations."""

    def _process_symbol_directive(self, call_node, directive):
        if len(call_node.args) < 1:
            raise ValueError(
                '"%s" requires a positional first argument' " as the target" % directive.__name__
            )
        target = call_node.args[0]
        defs = anno.getanno(target, anno.Static.ORIG_DEFINITIONS)
        for def_ in defs:
            def_.directives[directive] = _map_args(call_node, directive)
        return call_node

    def _process_statement_directive(self, call_node, directive):
        if self.state[_LoopScope].statements_visited > 1:
            raise ValueError(
                '"%s" must be the first statement in the loop block' % (directive.__name__)
            )
        if self.state[_LoopScope].level < 2:
            raise ValueError('"%s" must be used inside a statement' % directive.__name__)
        target = self.state[_LoopScope].ast_node
        node_anno = anno.getanno(target, anno.Basic.DIRECTIVES, {})
        node_anno[directive] = _map_args(call_node, directive)
        anno.setanno(target, anno.Basic.DIRECTIVES, node_anno)
        return call_node

    def visit_Name(self, node):
        node = self.generic_visit(node)
        if isinstance(node.ctx, gast.Load):
            defs = anno.getanno(node, anno.Static.DEFINITIONS, ())
            is_defined = bool(defs)
            if not is_defined and node.id in self.ctx.info.namespace:
                anno.setanno(node, STATIC_VALUE, self.ctx.info.namespace[node.id])
        return node

    def visit_Attribute(self, node):
        node = self.generic_visit(node)
        parent_val = anno.getanno(node.value, STATIC_VALUE, default=None)
        if parent_val is not None and inspect.ismodule(parent_val):
            if hasattr(parent_val, node.attr):
                anno.setanno(node, STATIC_VALUE, getattr(parent_val, node.attr))
        return node

    def visit_Assign(self, node):
        self.state[_LoopScope].statements_visited += 1
        return self.generic_visit(node)

    def visit_AugAssign(self, node):
        self.state[_LoopScope].statements_visited += 1
        return self.generic_visit(node)

    def visit_Expr(self, node):
        self.state[_LoopScope].statements_visited += 1
        node = self.generic_visit(node)
        if isinstance(node.value, gast.Call):
            call_node = node.value
            static_val = anno.getanno(call_node.func, STATIC_VALUE, default=None)
            if static_val is not None:
                # Note: directive calls are not output in the generated code, hence
                # the removal from the code by returning None.

                if static_val is directives.set_element_type:
                    self._process_symbol_directive(call_node, static_val)
                    return None
                elif static_val is directives.set_loop_options:
                    self._process_statement_directive(call_node, static_val)
                    return None
        return node

    # TODO(mdan): This will be insufficient for other control flow.
    # That means that if we ever have a directive that affects things other than
    # loops, we'll need support for parallel scopes, or have multiple converters.
    def _track_and_visit_loop(self, node):
        self.state[_LoopScope].enter()
        self.state[_LoopScope].ast_node = node
        node = self.generic_visit(node)
        # Edge case: a loop with just one directive statement would become empty.
        if not node.body:
            node.body = [gast.Pass()]
        self.state[_LoopScope].exit()
        return node

    def visit_While(self, node):
        return self._track_and_visit_loop(node)

    def visit_For(self, node):
        return self._track_and_visit_loop(node)


def transform(node, ctx):
    return DirectivesTransformer(ctx).visit(node)
